﻿/**************************************************************
 *
 * 唯一标识：ccd5256c-888f-4108-92ce-036a4cbad5c5
 * 命名空间：Sgr.Algorithm
 * 创建时间：2024/7/12 14:42:28
 * 机器名称：DESKTOP-HJ4OAG9
 * 创建者：CocoYuan
 * 电子邮箱：fengqinhua2016@163.com
 * 描述：
 *
 * 技术实现参考：
 * https://github.com/jlao/ConsistentHashing
 * https://github.com/tg123/ConsistentSharp
 **************************************************************/

using Sgr.Identity;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;

namespace Sgr.Algorithm
{
    public class DefaultConsistentHash : IConsistentHash
    {
        protected readonly ReaderWriterLockSlim _rwlock = new ReaderWriterLockSlim(LockRecursionPolicy.SupportsRecursion);
        protected readonly Dictionary<uint, string> _circle = new Dictionary<uint, string>();
        protected readonly Dictionary<string, uint[]> _nodes = new Dictionary<string, uint[]>();

        protected uint[] _sortedHashes = new uint[0];

        #region IConsistentHash

        /// <summary>
        /// 添加节点
        /// </summary>
        /// <param name="nodeName"></param>
        /// <param name="virtualNodeCount"></param>
        public virtual void AddNode(string nodeName, uint virtualNodeCount = 200)
        {
            Check.StringNotNullOrEmpty(nodeName, nameof(nodeName));
            if (virtualNodeCount < 1)
                virtualNodeCount = 1;

            _rwlock.EnterUpgradeableReadLock();
            try
            {
                if (!_nodes.ContainsKey(nodeName))
                {
                    _rwlock.EnterWriteLock();

                    try
                    {
                        uint[] uints = new uint[virtualNodeCount];
                        for (var i = 0; i < virtualNodeCount; i++)
                        {
                            uint key = HashKey(Encoding.UTF8.GetBytes($"{nodeName}:{i}"));
                            _circle.Add(key, nodeName);
                            uints[i] = key;
                        }
                        _nodes.Add(nodeName, uints);
                        UpdateSortedHashes();
                    }
                    finally
                    {
                        _rwlock.ExitWriteLock();
                    }
                }
            }
            finally
            {
                _rwlock.ExitUpgradeableReadLock();
            }
        }

        public virtual void RemoveNode(string nodeName)
        {
            Check.StringNotNullOrEmpty(nodeName, nameof(nodeName));

            _rwlock.EnterUpgradeableReadLock();

            try
            {
                if (_nodes.TryGetValue(nodeName, out uint[]? idxs))
                {
                    _rwlock.EnterWriteLock();
                    try
                    {
                        foreach (var i in idxs)
                        {
                            _circle.Remove(i);
                        }

                        _nodes.Remove(nodeName);
                        UpdateSortedHashes();
                    }
                    finally
                    {
                        _rwlock.ExitWriteLock();
                    }
                }
            }
            finally
            {
                _rwlock.ExitUpgradeableReadLock();
            }
        }

        public virtual IEnumerable<string> GetAllNodeNames()
        {
            _rwlock.EnterReadLock();

            try
            {
                return _nodes.Keys.ToArray();
            }
            finally
            {
                _rwlock.ExitReadLock();
            }
        }

        /// <summary>
        /// 获取所有Hash Key
        /// </summary>
        /// <returns></returns>
        public virtual IEnumerable<uint> GetAllHashKeys()
        {
            _rwlock.EnterReadLock();

            try
            {
                return _circle.Keys.ToArray();
            }
            finally
            {
                _rwlock.ExitReadLock();
            }
        }

        public virtual string SearchNodeName(ReadOnlySpan<byte> source)
        {
            _rwlock.EnterReadLock();

            try
            {
                if (_sortedHashes.Length == 0)
                    throw new EmptyCircleException("未设置节点");

                uint key = HashKey(source);
                var idx = BinarySearch((uint)_sortedHashes.Length, x => _sortedHashes[x] > key);

                if (idx >= _sortedHashes.Length)
                    idx = 0;
                return _circle[_sortedHashes[idx]];
            }
            finally
            {
                _rwlock.ExitReadLock();
            }
        }

        #endregion IConsistentHash

        #region IDisposable

        public virtual void Dispose()
        {
            _rwlock.Dispose();
        }

        #endregion IDisposable

        protected virtual uint BinarySearch(uint n, Func<uint, bool> f)
        {
            uint start = 0;
            uint end = n;

            while (start < end)
            {
                var mid = (start + end) / 2;

                if (!f(mid))
                    start = mid + 1;
                else
                    end = mid;
            }

            return start;
        }

        protected virtual void UpdateSortedHashes()
        {
            var hashes = _circle.Keys.ToArray();
            Array.Sort(hashes);
            _sortedHashes = hashes;
        }

        protected virtual uint HashKey(ReadOnlySpan<byte> source)
        {
            return System.IO.Hashing.Crc32.HashToUInt32(source);
        }
    }
}